package com.zh.note.leetcode.array;

import java.util.Arrays;
import java.util.PriorityQueue;

/**
 * 题目： 最小的k个数
 * 输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。
 * <p>
 * 示例 1：
 * <p>
 * 输入：arr = [3,2,1], k = 2
 * 输出：[1,2] 或者 [2,1]
 * 示例 2：
 * <p>
 * 输入：arr = [0,1,2,1], k = 1
 * 输出：[0]
 */
public class GetLeastNumbers {
    /**
     * 步骤：
     * 1.先把k个数放到大根堆里面
     * 2.遍历剩余元素，如果比堆顶小，则移除当前堆顶元素，把该元素加入堆
     * 3.遍历取出堆元素
     */
    public static int[] getLeastNumbers(int[] arr, int k) {

        int[] res = new int[k];
        if (k == 0) {
            return res;
        }
        // 大根堆
        PriorityQueue<Integer> queue = new PriorityQueue<>((o1, o2) -> o2 - o1);
        // 1、先将前k个元素 加入大根堆
        for (int i = 0; i < k; i++) {
            queue.offer(arr[i]);
        }
        // 2、比较k个之后的元素 与堆顶元素
        for (int i = k; i < arr.length; i++) {
            // peek只会查看堆顶，不会移除元素
            if (queue.peek() > arr[i]) {
                queue.poll();
                queue.offer(arr[i]);
            }
        }
        for (int i = 0; i < k; i++) {
            res[i] = queue.poll();
        }
        return res;

    }

    public static void main(String[] args) {
        int[] nums = {4, 5, 1, 6, 2, 7, 3};
        System.out.println(Arrays.toString(getLeastNumbers(nums, 4)));
    }
}
